﻿#include "WebSocketClientHandler.h"
#include "WebSocketClientAcceptor.h"
#include <sstream>
#include <string>
#include <ace/INet/HTTP_Request.h>
#include <ace/Codecs.h>
#include <ace/OS.h>
#include <stdint.h>

//sha1散列算法。网上copy过来的。
namespace sha1
{
    namespace // local
    {
        // Rotate an integer value to left.
        inline const unsigned int rol(const unsigned int value,
                const unsigned int steps)
        {
            return ((value << steps) | (value >> (32 - steps)));
        }

        // Sets the first 16 integers in the buffert to zero.
        // Used for clearing the W buffert.
        inline void clearWBuffert(unsigned int* buffert)
        {
            for (int pos = 16; --pos >= 0;)
            {
                buffert[pos] = 0;
            }
        }

        void innerHash(unsigned int* result, unsigned int* w)
        {
            unsigned int a = result[0];
            unsigned int b = result[1];
            unsigned int c = result[2];
            unsigned int d = result[3];
            unsigned int e = result[4];

            int round = 0;

            #define sha1macro(func,val) \
                        { \
                const unsigned int t = rol(a, 5) + (func) + e + val + w[round]; \
                                e = d; \
                                d = c; \
                                c = rol(b, 30); \
                                b = a; \
                                a = t; \
                        }

            while (round < 16)
            {
                sha1macro((b & c) | (~b & d), 0x5a827999)
                ++round;
            }
            while (round < 20)
            {
                w[round] = rol((w[round - 3] ^ w[round - 8] ^ w[round - 14] ^ w[round - 16]), 1);
                sha1macro((b & c) | (~b & d), 0x5a827999)
                ++round;
            }
            while (round < 40)
            {
                w[round] = rol((w[round - 3] ^ w[round - 8] ^ w[round - 14] ^ w[round - 16]), 1);
                sha1macro(b ^ c ^ d, 0x6ed9eba1)
                ++round;
            }
            while (round < 60)
            {
                w[round] = rol((w[round - 3] ^ w[round - 8] ^ w[round - 14] ^ w[round - 16]), 1);
                sha1macro((b & c) | (b & d) | (c & d), 0x8f1bbcdc)
                ++round;
            }
            while (round < 80)
            {
                w[round] = rol((w[round - 3] ^ w[round - 8] ^ w[round - 14] ^ w[round - 16]), 1);
                sha1macro(b ^ c ^ d, 0xca62c1d6)
                ++round;
            }

            #undef sha1macro

            result[0] += a;
            result[1] += b;
            result[2] += c;
            result[3] += d;
            result[4] += e;
        }
    } // namespace

    void calc(const void* src, const int bytelength, unsigned char* hash)
    {
        // Init the result array.
        unsigned int result[5] = { 0x67452301, 0xefcdab89, 0x98badcfe, 0x10325476, 0xc3d2e1f0 };

        // Cast the void src pointer to be the byte array we can work with.
        const unsigned char* sarray = (const unsigned char*) src;

        // The reusable round buffer
        unsigned int w[80];

        // Loop through all complete 64byte blocks.
        const int endOfFullBlocks = bytelength - 64;
        int endCurrentBlock;
        int currentBlock = 0;

        while (currentBlock <= endOfFullBlocks)
        {
            endCurrentBlock = currentBlock + 64;

            // Init the round buffer with the 64 byte block data.
            for (int roundPos = 0; currentBlock < endCurrentBlock; currentBlock += 4)
            {
                // This line will swap endian on big endian and keep endian on little endian.
                w[roundPos++] = (unsigned int) sarray[currentBlock + 3]
                        | (((unsigned int) sarray[currentBlock + 2]) << 8)
                        | (((unsigned int) sarray[currentBlock + 1]) << 16)
                        | (((unsigned int) sarray[currentBlock]) << 24);
            }
            innerHash(result, w);
        }

        // Handle the last and not full 64 byte block if existing.
        endCurrentBlock = bytelength - currentBlock;
        clearWBuffert(w);
        int lastBlockBytes = 0;
        for (;lastBlockBytes < endCurrentBlock; ++lastBlockBytes)
        {
            w[lastBlockBytes >> 2] |= (unsigned int) sarray[lastBlockBytes + currentBlock] << ((3 - (lastBlockBytes & 3)) << 3);
        }
        w[lastBlockBytes >> 2] |= 0x80 << ((3 - (lastBlockBytes & 3)) << 3);
        if (endCurrentBlock >= 56)
        {
            innerHash(result, w);
            clearWBuffert(w);
        }
        w[15] = bytelength << 3;
        innerHash(result, w);

        // Store hash in result pointer, and make sure we get in in the correct order on both endian models.
        for (int hashByte = 20; --hashByte >= 0;)
        {
            hash[hashByte] = (result[hashByte >> 2] >> (((3 - hashByte) & 0x3) << 3)) & 0xff;
        }
    }

    void toHexString(const unsigned char* hash, char* hexstring)
    {
        const char hexDigits[] = { "0123456789abcdef" };

        for (int hashByte = 20; --hashByte >= 0;)
        {
            hexstring[hashByte << 1] = hexDigits[(hash[hashByte] >> 4) & 0xf];
            hexstring[(hashByte << 1) + 1] = hexDigits[hash[hashByte] & 0xf];
        }
        hexstring[40] = 0;
    }
}


WebSocketClientHandler::WebSocketClientHandler (void):mb(128*1024)
{
	
}

void WebSocketClientHandler::methodDone(){
}

//事情已经在destroy里面做完了，所以这里什么都不做。
WebSocketClientHandler::~WebSocketClientHandler (void)
{
}


void
WebSocketClientHandler::destroy (void)
{
  //DONT_CALL的意思是不要再去调用handle_close，否则会死循环。
  this->reactor ()->remove_handler (this,
                                    ACE_Event_Handler::READ_MASK
                                    | ACE_Event_Handler::DONT_CALL);  
  delete this;
}


int
WebSocketClientHandler::open (void *void_acceptor)
{
  /** 对端的地址 */
  ACE_INET_Addr addr;

  if (this->peer ().get_remote_addr (addr) == -1)
    return -1;

  this->state=INIT;
  
  WebSocketClientAcceptor *acceptor = (WebSocketClientAcceptor *) void_acceptor;


  if (acceptor->thread_per_connection ())
    return this->activate (THR_DETACHED);

  //只有单线程模式，才会走到下面

  this->reactor (acceptor->reactor ());


  if (this->reactor ()->register_handler (this,
                                          ACE_Event_Handler::READ_MASK) == -1)
    ACE_ERROR_RETURN ((LM_ERROR,
                       "(%P|%t) can't register with reactor\n"),
                      -1);

  
  ACE_DEBUG ((LM_DEBUG,
              "(%P|%t) connected with %s\n", addr.get_host_addr ()));

  /* Always return zero on success.  */
  return 0;
}

int
WebSocketClientHandler::close(u_long flags)
{
  ACE_UNUSED_ARG (flags);

  this->destroy ();

  return 0;
}

/** handle_input 如果 return -1，会走到这里来。*/
int
WebSocketClientHandler::handle_close (ACE_HANDLE handle,
                              ACE_Reactor_Mask mask)
{
  ACE_UNUSED_ARG (handle);
  ACE_UNUSED_ARG (mask);

  this->destroy ();
  return 0;
}


/**  只有单线程模式才会走到这里来。thread-per-connection不会*/
int
WebSocketClientHandler::handle_input (ACE_HANDLE handle)
{
  ACE_UNUSED_ARG (handle);

  char buf[mybufSize];

  return this->process (buf, sizeof (buf));
}



/** 只有thread-per-connection模式才会走到这里来 */
int
WebSocketClientHandler::svc(void)
{
  char buf[mybufSize];

  // 线程的主消息循环
  while( true )
    {
     
      if (this->process(buf, sizeof (buf)) == -1)
        return -1;
    }

  return 0;
}


int WebSocketClientHandler::process (char *rdbuf,
                         int rdbuf_len)
{
#ifndef NDEBUG
	memset(rdbuf,0,rdbuf_len);
#endif
  ssize_t bytes_read;
  switch ( (bytes_read = this->peer ().recv (rdbuf, rdbuf_len)) )
    {
    case -1:
      ACE_ERROR_RETURN ((LM_ERROR,
                         "(%P|%t) %p bad read\n",
                         "client"),
                        -1);
    case 0:
      ACE_ERROR_RETURN ((LM_ERROR,
                         "(%P|%t) closing daemon (fd = %d)\n",
                         this->get_handle ()),
                        -1);
    default:			
		if(this->mb.space()< (size_t)bytes_read){
			ACE_ERROR_RETURN ((LM_ERROR,
                         "(%P|%t) data too large, closing peer\n",
                         this->get_handle ()),
                        -1);
		}
		this->mb.copy(rdbuf, bytes_read);
		if(this->state==INIT){
			this->state=HandShake;
			if(parseHandShake()<0)
				return -1;
		} else if(this->state==DataTransfer){
			if(parseData()<0)
				return -1;
		} else {
			//throw new std::runtime_error("wrong state");
			return -1;
		}
    }

  return 0;
}

int WebSocketClientHandler::parseHandShake(){
	std::string  str(this->mb.rd_ptr(),this->mb.rd_ptr()+this->mb.length());
	//找到连续两个换行符，就代表header收完了。
	if(str.find("\r\n\r\n")==std::string::npos)
		return 0;
	

	std::istringstream iss(str);	
	ACE::HTTP::Request req;
	if(!req.read(iss))
		return -1;
	if(req.get_method() != ACE::HTTP::Request::HTTP_GET)
		return -1;//unsupported method;
	//h.write(std::cout);
    ACE_CString key,value;
	key="Sec-WebSocket-Key";
	if(!req.get(key,value)){
		ACE_ERROR((LM_ERROR,"required header Sec-WebSocket-Key not found"));
		return -1;
	}
	value+="258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
	unsigned char hash[20];
	sha1::calc(value.c_str(),value.length(),hash);
	size_t outputlen;
	ACE_Byte* res=ACE_Base64::encode(hash,sizeof(hash),&outputlen);
	std::ostringstream oss;
	oss<<"HTTP/1.1 101 Switching Protocols\r\n"
		"Upgrade: websocket\r\n"
		"Connection: Upgrade\r\n"
		"Sec-WebSocket-Accept: "<<(const char*)res<<"\r\n";		
	this->mb.reset();
	std::string s=oss.str();
	this->mb.copy(s.c_str(),s.length());
	while(mb.length()>0){
			ssize_t sent=this->peer().send(mb.rd_ptr(),mb.length());
			if(sent<=0) 
				break; //TODO: 关闭连接
			mb.rd_ptr(sent);
		}
	this->mb.reset();
	this->state=DataTransfer;
	this->uri=req.get_URI();
	return 0;
}

int WebSocketClientHandler::parseData(){
	char* rdptr=NULL;
	while(this->mb.length()>0){
		rdptr=this->mb.rd_ptr();
		//每个包至少得有2个字节的head和4个字节的mask key。
		if(this->mb.length()<6)
			goto end;
		//读head
		const uint8_t head1=*(uint8_t*)this->mb.rd_ptr();
		this->mb.rd_ptr(1);
		const uint8_t head2=*(uint8_t*)this->mb.rd_ptr();
		this->mb.rd_ptr(1);

		const uint8_t opcode= head1 & ((1<<4)-1);	
		if(opcode!=0x1)
			return -1; //只支持text frame，收到其它的一概不理，直接关闭连接
		const bool isFin= head1 & 0x80;
		const bool isMask= head2 & 0x80; //must be true!
		uint64_t payloadLen= head2 &  ((1<<7)-1);
		if(payloadLen==126){
			if(this->mb.length()<2)
				goto end;
			payloadLen=ntohs(*(uint16_t*)this->mb.rd_ptr());
			this->mb.rd_ptr(2);
		} else if(payloadLen==127){
			if(this->mb.length()<8)
				goto end;
			//XXX: 有没有什么标准函数是把64位整数从网络序转成主机序的？
			payloadLen=ntohl(*(uint32_t*)this->mb.rd_ptr());
			payloadLen<<=32;
			payloadLen+=ntohl(*(uint32_t*)this->mb.rd_ptr());
			this->mb.rd_ptr(8);
		}
		//TODO: 检查payloadLen是否大于指定值。例如，32位平台上不能大于pow(2,32) -1。

		uint8_t key[4];
		static const size_t MaskKeyLen=4;
		if(this->mb.length()<MaskKeyLen)
			goto end;
		memcpy(key,this->mb.rd_ptr(),MaskKeyLen);
		this->mb.rd_ptr(MaskKeyLen);

		uint8_t* data=new uint8_t[(size_t)payloadLen+1];
		if(this->mb.length()<payloadLen)
			goto end;
		memcpy(data,this->mb.rd_ptr(),(size_t)payloadLen);
		this->mb.rd_ptr((size_t)payloadLen);
		
		for(int i=0;i!=payloadLen;++i){
			data[i]^=key[i%4];
		}
		data[payloadLen]='\0';
		clientMessageStream<<data;		
		rdptr=NULL;
		if(isFin) {
			processMessage(clientMessageStream.str());
			clientMessageStream.str("");
		}
	}	
end:
	if(rdptr!=NULL)
		this->mb.rd_ptr(rdptr);
	this->mb.crunch();
	return 0;
}

void WebSocketClientHandler::sendReplyMessage(const std::string& reply){
	if(reply.length()>0xFFFFFFFF)
		return;
	uint8_t* buffer=new uint8_t [16+reply.length()];
	uint8_t* wp=buffer;
	*wp++=0x81; //FIN && text frame
	if(reply.length()<126){
		*wp++=reply.length();
	} else if(reply.length()<=65535){
		uint16_t t=htons(reply.length());
		memcpy(wp,&t,sizeof(uint16_t));
		wp+=sizeof(uint16_t);
	} else {
		const uint64_t len=reply.length();		
		uint32_t t=0; //高位始终为0。不可能发4G以上的包
		memcpy(wp,&t,sizeof(uint32_t));
		wp+=sizeof(uint32_t);
		uint32_t t1=htonl(len& 0xFFFFFFFFL);
		memcpy(wp,&t1,sizeof(uint32_t));
		wp+=sizeof(uint32_t);
	}
	memcpy(wp,reply.c_str(),reply.length());
	wp+=reply.length();
		
	uint8_t* rp=buffer;
	do{
		ssize_t sent=this->peer().send(buffer,wp-buffer);
		if(sent<=0) {
			delete[] buffer;
			return;
		}
		rp+=sent;
	}while(rp<wp);	
	delete[] buffer;
}

void WebSocketClientHandler::processMessage(const std::string& req){
	ACE_DEBUG((LM_DEBUG,"received message %s\n",req.c_str()));	
	this->sendReplyMessage("hello");
}